﻿using System;
using System.Security.Cryptography;
using System.Text;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Rhino.Mocks;
using SjclHelpers.Codec;
using SjclHelpers.Misc;

namespace SjclHelpersTests.Misc
{
    [TestClass]
    public class PBKDF2HMACSHA256Tests
    {
        public byte[] ValidPassword
        {
            get { return new byte[] { 1 }; }
        }

        public byte[] ValidSalt
        {
            get { return new byte[] { 1, 2, 3, 4, 5, 6, 7, 8 }; }
        }

        public int ValidIterations
        {
            get { return 1; }
        }

        [TestMethod]
        public void UIntMaxSameAs2Pow32Minus1()
        {
            // Exercise SUT...
            var twoPow32Minus1 = (uint)Math.Pow(2, 32) - 1;

            // Verify SUT...
            Assert.AreEqual(uint.MaxValue, twoPow32Minus1);
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  Ctor...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        [TestMethod]
        [ExpectedException(typeof(ArgumentNullException))]
        public void Ctor_PasswordNull_ArgumentNullException()
        {
            // Fixture setup...
            byte[] password = null;
        
            // Exercise and verify SUT...
            new PBKDF2HMACSHA256(password, ValidSalt, ValidIterations);
        }
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentOutOfRangeException))]
        public void Ctor_Password0Bytes_ArgumentOutOfRangeException()
        {
            // Fixture setup...
            var password = new byte[0];
        
            // Exercise and verify SUT...
            new PBKDF2HMACSHA256(password, ValidSalt, ValidIterations);
        }
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentNullException))]
        public void Ctor_SaltNull_ArgumentNullException()
        {
            // Fixture setup...
            var password = new byte[] { 0 };
            byte[] salt = null;
        
            // Exercise and verify SUT...
            new PBKDF2HMACSHA256(ValidPassword, salt, ValidIterations);
        }
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentOutOfRangeException))]
        public void Ctor_SaltLessThan8Bytes_ArgumentOutOfRangeException()
        {
            // Fixture setup...
            var salt = new byte[] { 0, 0, 0, 0, 0, 0, 0 };
        
            // Exercise and verify SUT...
            new PBKDF2HMACSHA256(ValidPassword, salt, 1);
        }
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentOutOfRangeException))]
        public void Ctor_Iterations0_ArgumentOutOfRangeException()
        {
            // Fixture setup...
            var iterations = 0;
        
            // Exercise and verify SUT...
            new PBKDF2HMACSHA256(ValidPassword, ValidSalt, iterations);
        }
        
        [TestMethod]
        public void Ctor_ValidArgs_SetsPropsToArgs()
        {
            // Fixture setup...
        
            // Exercise SUT...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );
        
            // Verify SUT...
            Assert.IsNotNull(derivator.PseudorandomFunction);
            CollectionAssert.AreEqual(ValidSalt, derivator.Salt);
            Assert.AreEqual(ValidIterations, derivator.Iterations);
            Assert.AreEqual((uint)0, derivator.IndexOfNextBitOfKeyToFetch);
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  ConcatSaltAndBlockIndex...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentOutOfRangeException))]
        public void ConcatSaltAndBlockIndex_BlockIndex0_ArgumentOutOfRangeException()
        {
            // Fixture setup...
            var salt = new byte[] { 0, 0, 0, 0, 0, 0, 0, 0 };
            var derivator = new PBKDF2HMACSHA256(new byte[] { 0 }, salt, 1000);
            const int blockIndex = 0;
        
            // Exercise and verify SUT...
            derivator.ConcatSaltAndBlockIndex(blockIndex);
        }
        
        [TestMethod]
        public void ConcatSaltAndBlockIndex_WorksForMinBlockIndex()
        {
            // Fixture setup...
            var salt = new byte[] { 0, 0, 0, 0, 0, 0, 0, 1 };
            var derivator = new PBKDF2HMACSHA256(new byte[] { 0 }, salt, 1000);
            const int blockIndex = 1;
            var expectedBytes = new byte[12] {
                0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1
            };
        
            // Exercise SUT...
            var actualBytes = derivator.ConcatSaltAndBlockIndex(blockIndex);
        
            // Verify SUT...
            CollectionAssert.AreEqual(expectedBytes, actualBytes);
        }
        
        [TestMethod]
        public void ConcatSaltAndBlockIndex_WorksForMaxBlockIndex()
        {
            // Fixture setup...
            var salt = new byte[8] { 1, 0, 0, 0, 0, 0, 0, 0 };
            var derivator = new PBKDF2HMACSHA256(new byte[] { 0 }, salt, 1000);
            const uint blockIndex = uint.MaxValue;
            var expectedBytes = new byte[12] {
                1, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 255
            };
        
            // Exercise SUT...
            var actualBytes = derivator.ConcatSaltAndBlockIndex(blockIndex);
        
            // Verify SUT...
            CollectionAssert.AreEqual(expectedBytes, actualBytes);
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  Randomize...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        [TestMethod]
        public void Randomize_CallsFunc()
        {
            // Fixture setup...
            var password = Encoding.UTF8.GetBytes("secret");
            var block = "722ae2da49249d03e40cc9357c8e88c5".ToBytes();
            var expectedRandomizedBlock = new HMACSHA256(password).ComputeHash(block);
            var derivator = new PBKDF2HMACSHA256(
                    password, ValidSalt, ValidIterations
                );

            // Exercise SUT...
            var actualRandomizedBlock = derivator.Randomize(block);

            // Verify SUT ...
            CollectionAssert.AreEqual(expectedRandomizedBlock, actualRandomizedBlock);
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  ComputeBlock...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        
        [TestMethod]
        [ExpectedException(typeof(ArgumentNullException))]
        public void ComputeBlock_BlockNull_ArgumentNullException()
        {
            // Fixture setup...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );
        
            // Exercise and verify SUT...
            derivator.ComputeBlock(null);
        }

        [TestMethod]
        public void ComputeBlock_1Iteration_RandomizeOnce()
        {
            // Fixture setup...
            const int iterations = 1;
            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, iterations
                );
            var block = new byte [] { 1 };
            var randomizedBlock = new byte [] { 0 };
            mockDerivator
                .Expect(d => d.Randomize(block))
                .Return(randomizedBlock);
            mockDerivator
                .Expect(d => d.Iterations)
                .Return(iterations);
        
            // Exercise SUT...
            var computedBlock = mockDerivator.ComputeBlock(block);

            // Verify SUT ...
            mockDerivator.VerifyAllExpectations();
            Assert.AreEqual(randomizedBlock, computedBlock);        
        }

        [TestMethod]
        public void ComputeBlock_1Iteration_VerifiedVsSjcl()
        {
            // Fixture setup...
            var password = Encoding.UTF8.GetBytes("secret");
            var salt = "4ae2f9b5340f64a5".ToBytes();
            const int iterations = 1;
            var derivator = new PBKDF2HMACSHA256(
                password, salt, iterations
            );
            var block = derivator.ConcatSaltAndBlockIndex(1);

            // Exercise SUT...
            var computedBlock = derivator.ComputeBlock(block);

            // Verify SUT ...
            Assert.AreEqual(
                "6bd7dee090cc90252ea2dae78688d5c2c64405b87fe4f69e35c3261ab81ffc59", 
                computedBlock.ToHex()
            );
        }

        [TestMethod]
        public void ComputeBlock_2Iterations_RandomizeTwiceAndXor()
        {
            // Fixture setup...
            const int iterations = 2;
            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, iterations
                );
            var block = new byte[] { 0, 0 };
            var firstRandomizedBlock = new byte[] { 0, 1 };
            var secondRandomizedBlock = new byte[] { 1, 1 };
            mockDerivator
                .Expect(d => d.Randomize(block))
                .Return(firstRandomizedBlock);
            mockDerivator
                .Expect(d => d.Randomize(firstRandomizedBlock))
                .Return(secondRandomizedBlock);

            // Exercise SUT...
            var computedBlock = mockDerivator.ComputeBlock(block);

            // Verify SUT ...
            mockDerivator.VerifyAllExpectations();
            CollectionAssert.AreEqual(new byte [] { 1, 0, }, computedBlock);
        }

        [TestMethod]
        public void ComputeBlock_2Iterations_VerifiedVsSjcl()
        {
            // Fixture setup...
            var password = Encoding.UTF8.GetBytes("secret2");
            var salt = "57ed52d74cfc5c60".ToBytes();
            const int iterations = 2;
            var derivator = new PBKDF2HMACSHA256(
                password, salt, iterations
            );
            var block = derivator.ConcatSaltAndBlockIndex(1);

            // Exercise SUT...
            var computedBlock = derivator.ComputeBlock(block);

            // Verify SUT ...
            Assert.AreEqual(
                "078fac0284f373604a6c3a047fdb2e10e256e9ffb586dd324f2a69dd9a6c8451",
                computedBlock.ToHex()
            );
        }

        [TestMethod]
        public void ComputeBlock_3Iterations_3RandomizeAndXor()
        {
            // Fixture setup...
            const int iterations = 3;
            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, iterations
                );
            var block = new byte[] { 0, 0 };
            var firstRandomizedBlock = new byte[] { 0, 0, 1 };
            var secondRandomizedBlock = new byte[] { 0, 1, 1 };
            var thirdRandomizedBlock = new byte[] { 1, 1, 1 };
            mockDerivator
                .Expect(d => d.Randomize(block))
                .Return(firstRandomizedBlock);
            mockDerivator
                .Expect(d => d.Randomize(firstRandomizedBlock))
                .Return(secondRandomizedBlock);
            mockDerivator
                .Expect(d => d.Randomize(secondRandomizedBlock))
                .Return(thirdRandomizedBlock);

            // Exercise SUT...
            var computedBlock = mockDerivator.ComputeBlock(block);

            // Verify SUT ...
            mockDerivator.VerifyAllExpectations();
            CollectionAssert.AreEqual(new byte[] { 1, 0, 1 }, computedBlock);
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  GetBytes...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        [TestMethod]
        [ExpectedException(typeof(ArgumentOutOfRangeException))]
        public void GetBytes_KeyLengthLessThan1_ArgumentOutOfRangeException()
        {
            // Fixture setup...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );
        
            // Exercise and verify SUT...
            derivator.GetBytes(0);
        }

        [TestMethod]
        [ExpectedException(typeof(NotImplementedException))]
        public void GetBytes_KeyLengthGreaterThan256_NotImplementedException()
        {
            // Fixture setup...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );

            // Exercise and verify SUT...
            derivator.GetBytes(257);
        }

        [TestMethod]
        [ExpectedException(typeof(NotImplementedException))]
        public void GetBytes_TotalKeyLengthGreaterThan256_NotImplementedException()
        {
            // Fixture setup...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );
            derivator.IndexOfNextBitOfKeyToFetch = 256;

            // Exercise and verify SUT...
            derivator.GetBytes(1);
        }

        [TestMethod]
        public void GetBytes_CallsFuncs()
        {
            // Fixture setup...
            const int numBitsToFetch = 1;
            var expectedBlockToCompute = new byte[] { };
            var expectedComputedBlock = new byte[] { 1, 0 };

            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, ValidIterations
                );
            mockDerivator
                .Expect(d => d.IndexOfNextBitOfKeyToFetch = (uint)numBitsToFetch);
            mockDerivator
                .Expect(d => d.ConcatSaltAndBlockIndex(1))
                .Return(expectedBlockToCompute);
            mockDerivator
                .Expect(d => d.ComputeBlock(expectedBlockToCompute))
                .Return(expectedComputedBlock);
            mockDerivator
                .Expect(d => d.LastComputedBlock = expectedComputedBlock);

            // Exercise SUT...
            var bytes = mockDerivator.GetBytes(numBitsToFetch);

            // Verify SUT...
            mockDerivator.VerifyAllExpectations();
            CollectionAssert.AreEqual(
                new byte[] { expectedComputedBlock[0] }, bytes
            );
        }

        [TestMethod]
        public void GetBytes_SecondCallNoReset_CallsFuncs()
        {
            // Fixture setup...
            const uint indexOfNextBitOfKeyToFetch = 1;
            const int numBitsToFetch = 1;
            var lastComputedBlock = new byte[] { 1, 0 };

            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, ValidIterations
                );
            mockDerivator
                .Expect(d => d.IndexOfNextBitOfKeyToFetch)
                .Return(indexOfNextBitOfKeyToFetch);
            mockDerivator
                .Expect(d => d.LastComputedBlock)
                .Return(lastComputedBlock)
                .Repeat.Twice();
            mockDerivator
                .Expect(
                    d => d.IndexOfNextBitOfKeyToFetch =
                        indexOfNextBitOfKeyToFetch + numBitsToFetch
                );

            // Exercise SUT...
            var bytes = mockDerivator.GetBytes(numBitsToFetch);

            // Verify SUT...
            mockDerivator.VerifyAllExpectations();
            CollectionAssert.AreEqual(
                new byte[] { lastComputedBlock[1] }, bytes
            );
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void GetBytes_IndexOfNextSetNoLastBlock_InvalidOperationException()
        {
            // Fixture setup...
            var derivator = new PBKDF2HMACSHA256(
                ValidPassword, ValidSalt, ValidIterations
            );
            derivator.IndexOfNextBitOfKeyToFetch = 1;
            derivator.LastComputedBlock = null;

            // Exercise and verify SUT...
            derivator.GetBytes(1);
        }

        [TestMethod]
        public void GetBytes_VerifiedVsSjcl()
        {
            // Fixture setup...
            var password = Encoding.UTF8.GetBytes("a super secret password");
            var salt = "b7678377db310ad6".ToBytes();
            const int iterations = 1000;
            const int keyLength = 256;
            var derivator = new PBKDF2HMACSHA256(
                password, salt, iterations
            );

            // Exercise SUT...
            var derivedKey = derivator.GetBytes(keyLength);

            // Verify SUT...
            Assert.AreEqual(
                "c93cda460fe7e990957ae7ff2f58d0590da346b54de1ee41672672998b973a0d",
                derivedKey.ToHex()
            );
        }

        [TestMethod]
        public void GetBytes_CalledTwice_VerifiedVsSjcl()
        {
            // Fixture setup...
            var password = Encoding.UTF8.GetBytes("pwd");
            var salt = "d54c8d0fc33dff75".ToBytes();
            const int iterations = 500;
            const int keyLength = 64 / 8;
            var derivator = new PBKDF2HMACSHA256(
                password, salt, iterations
            );

            // Exercise SUT...
            var derivedKeyPart1 = derivator.GetBytes(keyLength);
            var derivedKeyPart2 = derivator.GetBytes(keyLength);

            // Verify SUT...
            Assert.AreEqual(
                "0382f43f952b66b2762b6eaae2e23728",
                derivedKeyPart1.ToHex() + derivedKeyPart2.ToHex()
            );
        }

        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        //  Reset...
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        [TestMethod]
        public void Reset_SetsProps()
        {
            // Fixture setup...
            var mockDerivator = MockRepository
                .GeneratePartialMock<PBKDF2HMACSHA256>(
                    ValidPassword, ValidSalt, ValidIterations
                );
            mockDerivator
                .Expect(d => d.IndexOfNextBitOfKeyToFetch = 0);
            mockDerivator
                .Expect(d => d.LastComputedBlock = null);

            // Exercise SUT...
            mockDerivator.Reset();

            // Verify SUT...
            mockDerivator.VerifyAllExpectations();
        }
    }
}
